"""
阿里百炼DashScope向量模型实现
"""

import os
from typing import List

from openai import OpenAI

from config.logging_config import get_logger
from .BaseEmb import BaseEmb, ConfigurationError, EmbeddingError

logger = get_logger(__name__)


class DashScopeEmb(BaseEmb):
    """阿里百炼DashScope向量模型实现"""

    # 支持的模型列表
    SUPPORTED_MODELS = [
        "text-embedding-v3",
        "text-embedding-v4"
    ]

    # 默认配置
    DEFAULT_DIMENSIONS = {
        "text-embedding-v3": 1024,
        "text-embedding-v4": 1024
    }

    def __init__(self, model_name: str, api_key: str = None, dimensions: int = None, **kwargs):
        """
        初始化阿里百炼DashScope向量模型
        
        Args:
            model_name: 模型名称
            api_key: API密钥，如果未提供则从环境变量DASHSCOPE_API_KEY获取
            dimensions: 向量维度（仅v3和v4支持）
            **kwargs: 其他配置参数
        """
        self.api_key = api_key or os.getenv("DASHSCOPE_API_KEY")
        self.dimensions = dimensions
        super().__init__(model_name, **kwargs)

        # 初始化客户端
        try:
            self.client = OpenAI(
                api_key=self.api_key,
                base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
            )
        except Exception as e:
            raise ConfigurationError(f"初始化DashScope客户端失败: {str(e)}")

    def _validate_config(self) -> None:
        """验证配置参数"""
        if not self.api_key:
            raise ConfigurationError("DashScope API密钥不能为空，请设置DASHSCOPE_API_KEY环境变量或传入api_key参数")

        if self.model_name not in self.SUPPORTED_MODELS:
            logger.warning(f"模型 {self.model_name} 不在支持列表中，可能无法正常工作")

        # 验证dimensions参数
        if self.dimensions is not None:
            if self.model_name not in ["text-embedding-v3", "text-embedding-v4"]:
                logger.warning(f"模型 {self.model_name} 不支持自定义维度参数")
            elif self.dimensions <= 0:
                raise ConfigurationError("向量维度必须大于0")

    def get_emb(self, text: str) -> List[float]:
        """
        获取单个文本的向量表示
        
        Args:
            text: 输入文本
            
        Returns:
            向量列表
            
        Raises:
            EmbeddingError: 向量化失败时抛出
        """
        if not text or not text.strip():
            raise ValueError("输入文本不能为空")

        try:
            # 准备请求参数
            params = {
                "model": self.model_name,
                "input": text.strip(),
                "encoding_format": "float"
            }

            # 对于支持dimensions的模型添加该参数
            if self.dimensions and self.model_name in ["text-embedding-v3", "text-embedding-v4"]:
                params["dimensions"] = self.dimensions

            response = self.client.embeddings.create(**params)

            if not response.data or len(response.data) == 0:
                raise EmbeddingError("API返回空结果")

            return response.data[0].embedding

        except Exception as e:
            if isinstance(e, (ValueError, EmbeddingError)):
                raise
            logger.error(f"DashScope向量化失败: {str(e)}")
            raise EmbeddingError(f"向量化失败: {str(e)}")

    def encode(self, texts: List[str]) -> List[List[float]]:
        """
        编码多个文本
        
        Args:
            texts: 文本列表
            
        Returns:
            向量列表的列表
        """
        if not texts:
            return []

        # 过滤空文本
        valid_texts = [text.strip() for text in texts if text and text.strip()]
        if not valid_texts:
            raise ValueError("没有有效的输入文本")

        try:
            # 准备请求参数
            params = {
                "model": self.model_name,
                "input": valid_texts,
                "encoding_format": "float"
            }

            # 对于支持dimensions的模型添加该参数
            if self.dimensions and self.model_name in ["text-embedding-v3", "text-embedding-v4"]:
                params["dimensions"] = self.dimensions

            response = self.client.embeddings.create(**params)

            if not response.data or len(response.data) != len(valid_texts):
                raise EmbeddingError("API返回结果数量不匹配")

            # 按顺序返回结果
            results = []
            for item in response.data:
                results.append(item.embedding)

            return results

        except Exception as e:
            if isinstance(e, (ValueError, EmbeddingError)):
                raise
            logger.error(f"DashScope批量向量化失败: {str(e)}")
            raise EmbeddingError(f"批量向量化失败: {str(e)}")

    def encode_batch(self, texts: List[str], batch_size: int = 50) -> List[List[float]]:
        """
        批量编码文本（DashScope优化版本）
        
        Args:
            texts: 文本列表
            batch_size: 批次大小，建议不超过50
            
        Returns:
            向量列表的列表
        """
        if not texts:
            return []

        # DashScope建议批次大小不超过50
        batch_size = min(batch_size, 50)

        if len(texts) <= batch_size:
            return self.encode(texts)

        results = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            try:
                batch_results = self.encode(batch)
                results.extend(batch_results)
            except Exception as e:
                logger.error(f"批次 {i // batch_size + 1} 处理失败: {str(e)}")
                raise

        return results

    def get_dimension(self) -> int:
        """
        获取向量维度
        
        Returns:
            向量维度
        """
        # 如果指定了维度，直接返回
        if self.dimensions:
            return self.dimensions

        # 如果有默认维度配置，使用默认值
        if self.model_name in self.DEFAULT_DIMENSIONS:
            return self.DEFAULT_DIMENSIONS[self.model_name]

        # 否则通过实际调用获取
        return super().get_dimension()

    def get_model_info(self) -> dict:
        """
        获取模型信息
        
        Returns:
            模型信息字典
        """
        return {
            "provider": "Alibaba DashScope",
            "model_name": self.model_name,
            "supported_models": self.SUPPORTED_MODELS,
            "max_batch_size": 50,
            "dimension": self.get_dimension(),
            "supports_custom_dimensions": self.model_name in ["text-embedding-v3", "text-embedding-v4"],
            "current_dimensions": self.dimensions
        }

    def test_connection(self) -> bool:
        """
        测试连接是否正常
        
        Returns:
            连接是否正常
        """
        try:
            test_vector = self.get_emb("测试连接")
            return len(test_vector) > 0
        except Exception as e:
            logger.error(f"DashScope连接测试失败: {str(e)}")
            return False
